import gradio as gr
import numpy as np
from PIL import Image
from pinecone import Pinecone
from collections import Counter

# =================== 配置 ===================
PINECONE_API_KEY = "pcsk_7GvxrM_SVD5dH79WLzgGWd4JXQsdCoFHaXxsE3mJD83C4UXN8Ze1rRfkNC3ecQapfcd9DT"
INDEX_NAME = "mnist-knn-index"  # 确保这个索引已存在且包含训练数据
TOP_K = 11  # 使用 11 个最近邻做投票

# =================== 初始化 Pinecone ===================
print("Initializing Pinecone...")
pc = Pinecone(api_key=PINECONE_API_KEY)

if INDEX_NAME not in pc.list_indexes().names():
    raise ValueError(f"❌ 索引 '{INDEX_NAME}' 不存在，请先创建并上传 MNIST 训练数据")

index = pc.Index(INDEX_NAME)
print(f"✅ 成功连接到 Pinecone 索引: {INDEX_NAME}")

# =================== 预测函数 ===================
def predict_digit(drawing):
    try:
        # 获取绘制的图像（Gradio Sketchpad 返回的是字典）
        image = drawing['composite']  # 获取合成图像

        # 调整为 8x8 像素（与 MNIST 数据集一致）
        image = image.resize((8, 8), Image.Resampling.LANCZOS)

        # 转为灰度图
        image = image.convert('L')

        # 转为 numpy 数组
        image_array = np.array(image)  # 范围 0~255

        # 归一化到 0~15（sklearn 的 load_digits() 范围是 0~16，接近 0~15）
        image_array = (image_array / 255.0) * 15
        image_array = image_array.astype(np.float32).flatten()  # 展平为 64 维

        # =================== 发送到 Pinecone 查询 ===================
        results = index.query(
            vector=image_array.tolist(),  # 必须转为 list
            top_k=TOP_K,
            include_metadata=True
        )

        # 提取标签
        labels = []
        for match in results['matches']:
            if 'metadata' in match and 'label' in match['metadata']:
                labels.append(match['metadata']['label'])

        # 判断是否找到标签
        if not labels:
            return "❓ 未找到匹配项"

        # 多数投票
        predicted_label = Counter(labels).most_common(1)[0][0]

        return str(predicted_label)

    except Exception as e:
        return f"❌ 推理出错: {str(e)}"

# =================== 创建 Gradio 界面 ===================
def create_gradio_interface():
    # 兼容旧版 Gradio：去掉 brush_radius 和 canvas_size
    inp = gr.Sketchpad(label="Draw a digit", type="pil")

    iface = gr.Interface(
        fn=predict_digit,
        inputs=inp,
        outputs=gr.Textbox(label="Predicted Digit"),
        title="🎨 Handwritten Digit Recognizer (Powered by Pinecone)",
        description="Draw a digit (0-9) and see the prediction using vector similarity search in Pinecone!",
        live=False
    )

    return iface

# =================== 主程序 ===================
if __name__ == "__main__":
    demo = create_gradio_interface()
    print("Launching Gradio interface...")
    demo.launch(share=True)  # share=True 可生成公网链接